import os
import numpy as np
import pandas as pd
import zipfile
from utils.metrics import cal_csv_mAP
from sklearn.model_selection import train_test_split
from FSdata.FSdataset import FSdata, collate_fn,attr2length_map, idx2attr_map,attr2idx_map
from utils.preprocessing import join_path_to_df

def str2np(x):
    return np.array(x.split(';')).astype(np.float)

def np2str(arr):
    return ';'.join(['%.4f'%x for x in arr])

rawdata_root = '/media/gserver/data/FashionAI'
round2_df = pd.read_csv(os.path.join(rawdata_root,'round2/train/Annotations/label.csv'),
                        header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
round2_df = join_path_to_df(round2_df, rawdata_root, 'round2/train')

round2_train_pd, val_pd = train_test_split(round2_df, test_size=0.1, random_state=37,
                                    stratify=round2_df['AttrKey'])


data_set_val = FSdata(
                     anno_pd=val_pd,
                     transforms=None,
                         )
mode = 'offline'

model_name = 'val-merge'
merge_list = [
    "res101_d[r2-4]9748-aug14-[0.9798].csv",
]

if mode == 'online':
    csv_root = './online_pred/csv/'
else:
    csv_root = './val_pred2/part_csv/'

val_csv =val_pd.drop('AttrValues',axis=1)

file_path = os.path.join(csv_root, merge_list[0])
merged_pd = pd.read_csv(file_path,header=None, names=['ImageName', 'AttrKey', 'AttrValueProbs'])
# merged_pd = pd.merge(val_csv, merged_pd, on=['ImageName','AttrKey'], how="left")
merged_pd['AttrValueProbs'] = merged_pd['AttrValueProbs'].apply(str2np)
print(merged_pd.keys())

for file_name in merge_list[1:]:
    file_path = os.path.join(csv_root, file_name)
    part_pd = pd.read_csv(file_path,header=None, names=['ImageName', 'AttrKey', 'AttrValueProbs'])
    # part_pd = pd.merge(val_csv, part_pd,on=['ImageName','AttrKey'],how="left")
    merged_pd['AttrValueProbs'] += part_pd['AttrValueProbs'].apply(str2np)

merged_pd['AttrValueProbs'] = merged_pd['AttrValueProbs'] / len(merge_list)
print(val_pd.iloc[100],merged_pd.iloc[100])

val_mAP,APs, accs = cal_csv_mAP(merged_pd,val_pd, ['neckline_design_labels'])

print('=='*20)
print('val-mAP: %.4f'% (val_mAP))
for key in APs.keys():
    print('acc: %.4f, AP: %.4f %s' % (accs[key], APs[key], key))


# merged_pd['AttrValueProbs'] = merged_pd['AttrValueProbs'].apply(np2str)

# # make zip file
# if mode == 'online':
#     merged_pd[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred/csv/%s.csv' % model_name, header=None,
#                                                                index=False)
#     z = zipfile.ZipFile('./online_pred/subs_zip/%s.zip'%model_name, 'w', zipfile.ZIP_DEFLATED)
#     z.write('./online_pred/csv/%s.csv' % model_name, arcname='%s.csv' % model_name)
#     z.close()
#
# else:
#     merged_pd[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('val_pred/csv/%s.csv' % model_name, header=None,
#                                                                index=False)
